import tensorflow as tf
from tensorflow import keras
import os
current_path = os.path.dirname(os.path.abspath(__file__))

(train_iamges,train_labels),(test_images,test_labels) = tf.keras.datasets.mnist.load_data()
train_labels = train_labels[:1000]
test_labels = test_labels[:1000]

train_iamges = train_iamges[:1000].reshape(-1,28*28)/255.0
test_images = test_images[:1000].reshape(-1,28*28)/255.0

def create_model():
    model = tf.keras.models.Sequential([
        keras.layers.Dense(512,activation='relu',input_shape=(784,)),
        keras.layers.Dropout(0.2),
        keras.layers.Dense(10,activation='softmax')
    ])
    model.compile(optimizer='adam',
                  loss="sparse_categorical_crossentropy",
                  metrics=['accuracy'])
    return model

# 创建基本模型实例
model = create_model()
model.summary()

checkpoint_path = os.path.join(current_path,"training_1/cp.ckpt")
checkpoint_dir = os.path.dirname(checkpoint_path)

# 创建一个检查点回调
cp_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1,)

model.fit(train_iamges,train_labels,epochs=10,
          validation_data=(test_images,test_labels),
          callbacks=[cp_callback])

# 这可能会生成与保存优化程序状态相关的警告。
# 这些警告（以及整个笔记本中的类似警告）是为了阻止过时使用的，可以忽略。

model =  create_model()
loss,acc = model.evaluate(test_images,test_labels)
print("Untrained model, accuracy: {:5.2f}%".format(100*acc))

model.load_weights(checkpoint_path)
loss,acc= model.evaluate(test_images,test_labels)
print("Restored model, accuracy: {:5.2f}%".format(100*acc))



#Keras使用HDF5标准提供基本保存格式，出于我们的目的，可以将保存的模型视为单个二进制blob。
model = create_model()

model.fit(train_iamges, train_labels, epochs=5)

# 保存整个模型到HDF5文件
model.save('my_model.h5')
loss, acc = model.evaluate(test_images, test_labels)
print("before Restored model, accuracy: {:5.2f}%".format(100*acc))
new_model = keras.models.load_model("my_model.h5")
new_model.summary()
loss, acc = new_model.evaluate(test_images, test_labels)
print("Restored model, accuracy: {:5.2f}%".format(100*acc))